#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <dirent.h>
#include <unistd.h>
#include <sys/inotify.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <signal.h>
#include <string>
#include <queue>
#include <vector>
#include <sstream>
#include <map>
#include <sys/socket.h>
#include <netinet/in.h>
#include <pthread.h>
#include <algorithm>

#define EVENT_SIZE ( sizeof (struct inotify_event) )

void *handle_usb_events(void*);
void *handle_network_events(void*);

bool add_device(std::string devname, std::vector<std::string> &results);
bool IsBadCharacter(char c);
bool ismine(std::string cmd);

void add_command(std::string cmd);
std::string get_command();
int get_command_count();

void handle_error (int error);
void get_event(int fd, const char *target);
void remove_device(std::string devname, std::vector<std::string> &results);
void Split(std::string str, std::string delim, std::vector<std::string> &results);
void get_usb_paths(std::vector<std::string> &results);
void kill_child(int child_pid);
void handle_command(std::string cmd);
void set_quitflag();
void *doprocessing(void *sock);
std::string run_command(std::string cmd);
std::string get_root_device();
std::string get_dev_label(std::string dlabel);

// Variables
std::string root_device;
std::queue<std::string> commands;
std::vector<std::string> drive_paths;
std::vector<std::string> drive_labels;
std::vector<std::string> drive_ids;
std::map<std::string, int> pidlink;

pthread_mutex_t threadmutex = PTHREAD_MUTEX_INITIALIZER;
bool QuitFlag = false;
bool DebugFlag = false;

using namespace std;


int main(int argc, char **argv)
{
	if (argc > 1) { 
		string oline = string(*++argv); // Eeewww. I know, I'm sorry.
		if (oline == "--debug") DebugFlag = true;
	}

	// Start The Daemonization Process
	pid_t this_pid = fork();
	if (this_pid < 0) exit(1);
	if (this_pid > 0) exit(0);


	// If We Are The Child (Now Main) Pid, Set rwx
	umask(0);


	// Become An Emancipated Minor
	pid_t sid;
	sid = setsid();
	if (sid < 0) {
		// log error
		exit(1);
	}


	// We Do This To Make Sure We Aren't On A Mount Point
	// If They Run A Daemon From A Mounted Device, They
	// Won't Be Able To Unmount The Device While Our Daemon
	// Is Running
	if ((chdir("/")) < 0) {
		// log error
		exit(1);
	}


	// Create Threads To Handle USB And Network Events
	int rc1, rc2;
	pthread_t usbthread, networkthread;

	if ((rc1=pthread_create(&usbthread, NULL, &handle_usb_events, NULL))) {
		cout << "Thead Creation Failed: " << rc1 << endl;
	}

	if ((rc2=pthread_create(&networkthread, NULL, &handle_network_events, NULL))) {
		cout << "Thead Creation Failed: " << rc2 << endl;
	}


	// We're A Daemon, We Don't Need These
	if (! DebugFlag) {
		close(STDIN_FILENO);
		close(STDOUT_FILENO);
		close(STDERR_FILENO);
	}
	
	while (!QuitFlag) {
		sleep(2);
	}

	

	return 0;
		
}

void *handle_usb_events(void*)
{
	drive_paths.clear();

	// Prep Inotify For Use
	const std::string directory = "/dev";
	int fd = inotify_init();
	inotify_add_watch(fd, directory.c_str(), IN_CREATE | IN_DELETE);
	int i = 0;
	char buf[1024 * (sizeof(inotify_event) + 16)];
	ssize_t length;
	std::string dlabel;


	// Get The Device That / Is Mounted on
	// Also Get All Of the Existing /dev/sd?? Devices
	root_device = get_root_device();
	get_usb_paths(drive_paths);


	// Main USB Loop
	while (1) {
		
		length = read(fd, buf, sizeof(buf));
		i = 0;

		while (i < length) {
			struct inotify_event *event = (struct inotify_event *) &buf[i];
			if (event->len) {
				if (event->mask & IN_CREATE) {
					// Note The Following Line: Throwing a NOT (!) in the mix made it act up
					// So I left the initial condition block blank, and stuck the else on the
					// same Line.
					if (event->mask & IN_ISDIR) { } else {
						std::string devname(event->name);
						if (add_device(devname, drive_paths)) {
							dlabel = get_dev_label(devname);
							pid_t pID = fork();

							if (pID == 0) { // child
								// child stuff
							} else {
								pidlink[devname] = pID;
							}
						}
					}
				} else if (event->mask & IN_DELETE) {
					// Note The Following Line: Throwing a NOT (!) in the mix made it act up
					// So I left the initial condition block blank, and stuck the else on the
					// same Line.
					if (event->mask & IN_ISDIR) { } else {
						std::string devname(event->name);
						std::map<std::string,int>::iterator it;

						it = pidlink.find(devname);
						if (it != pidlink.end()) {
							kill_child(pidlink[devname]);
							pidlink.erase(it);
						}
						
						remove_device(devname, drive_paths);
					}
				}
			}
	
			i += EVENT_SIZE + event->len;
		}
	}
	return 0;

}

void *handle_network_events(void*)
{
	// Socket Variables 
	int sockfd, newsockfd, portno, clilen;
	char buffer[256];
	struct sockaddr_in serv_addr, cli_addr;
	int  n;

	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd < 0) {
		perror("ERROR opening socket");
		exit(1);
	}

	// Initialize socket structure
	memset((char *) &serv_addr, 0, sizeof(serv_addr)); 
	portno = 5001;
	serv_addr.sin_family = AF_INET;
	serv_addr.sin_addr.s_addr = INADDR_ANY;
	serv_addr.sin_port = htons(portno);
 
	if (bind(sockfd, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0) {
		perror("ERROR on binding");
		exit(1);
	}

	listen(sockfd,5);
	clilen = sizeof(cli_addr);

	// Main Network Loop
	while (1) {
		newsockfd = accept(sockfd, (struct sockaddr *) &cli_addr, (socklen_t *) &clilen);

		if (newsockfd < 0) {
			perror("ERROR on accept");
			exit(1);
        	}
		
		int t1;
		pthread_t sockthread;

		if ((t1=pthread_create(&sockthread, NULL, &doprocessing, (void *) &newsockfd))) {
			exit(1);
		}
	}
}

bool IsBadCharacter(char c)
{
	switch (c) {
		case '\n':
		case '\r':
		case '\t':
		case '\0':
		case ' ':
			return true;
		default:
			return false;
	}
}

void *doprocessing(void *sock) 
{
	int n;
	char buffer[256];
	int client;

	client = *((int *)sock);
	memset(buffer, 0, 256);
	string cmd;

	while (cmd != "netquit") {
		memset(buffer, 0, 256);
		n = read(client, buffer, 255);

		if (n >= 0) {
			cmd = string(buffer);
			cmd.erase(std::remove_if(cmd.begin(), cmd.end(), &IsBadCharacter), cmd.end());
			handle_command(cmd);

		} else {
			close(client);
			pthread_exit(0);
		}
	}

	close(client);
	pthread_exit(0);
	return NULL;
}

bool ismine(std::string cmd)
{
	std::vector<std::string> dplist;
	dplist.push_back("netquit");

	bool flag = false;
	for (int i=0; i < dplist.size(); i++) if (cmd == dplist[i]) flag = true;
	return flag;
}

void handle_command(std::string cmd)
{
	std::map<std::string, void(*)()> cmdfunctions;
	cmdfunctions["shutdown"] = &set_quitflag;

	cmdfunctions[cmd]();
}

void set_quitflag()
{
	pthread_mutex_lock(&threadmutex);
	QuitFlag = true;
	pthread_mutex_unlock(&threadmutex);
}

void add_command(string cmd)
{
	pthread_mutex_lock(&threadmutex);
	commands.push(cmd);
	pthread_mutex_unlock(&threadmutex);
}

std::string peek_command()
{
	std::string tmp;
	pthread_mutex_lock(&threadmutex);
	if (commands.empty()) return "null";
	if (commands.size() <= 0 ) return "null";
	tmp = commands.front();
	pthread_mutex_unlock(&threadmutex);

	return tmp;
}

std::string get_command()
{
	std::string tmp;
	pthread_mutex_lock(&threadmutex);
	if (commands.empty()) return "null";
	if (commands.size() <= 0 ) return "null";
	tmp = commands.front();
	commands.pop();
	pthread_mutex_unlock(&threadmutex);

	return tmp;
}

int get_command_count()
{
	int cnt;
	pthread_mutex_lock(&threadmutex);
	cnt = commands.size();
	pthread_mutex_unlock(&threadmutex);
	return cnt;
}

std::string get_dev_label(std::string devname)
{
	sleep(4);
	std::vector<string> result;
	std::vector<string> pathparts;
	DIR *dir;
	char buf[1024];
	ssize_t len;
	struct dirent *ent;

	dir = opendir("/dev/disk/by-label");

	if (dir != NULL) {
		while (true) {
			ent = readdir(dir);
			if (ent == NULL) break;
			result.push_back(std::string(ent->d_name));
		}
	}

	for (size_t v=0; v < result.size(); v++) {
		if (result[v] == ".") result.erase(result.begin() + v);
		if (result[v] == "..") result.erase(result.begin() + v);
	}

	for (size_t v=0; v < result.size(); v++) {
		memset(buf, 0, sizeof(buf));
		std::string path = "/dev/disk/by-label/" + result[v];
		if ((len = readlink(path.c_str(), buf, sizeof(buf) -1)) != -1) buf[len] = '\0';

		Split((std::string)buf, "/", pathparts);

		if (pathparts[pathparts.size() - 1] == devname) return result[v];
	}

	return "null";
}

void kill_child(int child_pid)
{
	kill(child_pid, SIGKILL);
}

/* **************************************** */
/* Function To Split A String Into A Vector */
/* Based On A Delimiting Character In The   */
/* Passed In String                         */
/* **************************************** */
void Split(string str, string delim, vector<string> &results)
{
        std::string tmpresult;
        results.clear();

        // Loop For Each Character In The String 
        for (int i=0; i < (int)str.size(); i++) {
                if (str[i] == delim[0]) {

                        // The String Matches The Delimeter
                        results.push_back(tmpresult);
                        tmpresult.clear();

                } else {

                        tmpresult += str[i];

                }
        }

        // Stick The Remaining String Into The Vector
        results.push_back(tmpresult);
}

bool add_device(std::string devname, std::vector<std::string> &results)
{
	bool flag = false;
	bool cflag = false;
	bool aflag = false;
	
	for (size_t i=0; i < results.size(); i++) { 
		if (results[i] == devname) flag = true;
	}

	if (!flag) {
		int rchari;
		std::string rchar = devname.substr(devname.length() -1, devname.length());
		std::stringstream ss(rchar);
		ss >> rchari;

		for (int r = 1; (int)r < 10; r++) if (r == rchari) cflag = true;
		
		if (cflag) { 
			std::string ft = devname.substr(0, 2);
			if (ft == "sd") {
				results.push_back(devname);
				aflag = true;
			}
		}
	}

	return aflag;
}

void remove_device(std::string devname, std::vector<std::string> &results)
{
	for (size_t i=0; i < results.size(); i++) if (results[i] == devname) { results.erase(results.begin() + i); }
	return;
}

void get_usb_paths(std::vector<std::string> &results)
{
	std::string lsdata;
	std::string tmpdev;
	std::string rdev;
	std::vector<string> lineparts;
	size_t found;

	rdev = root_device.substr(0, root_device.size() - 1);
	lsdata = run_command("ls -1 /dev/sd??");
	Split(lsdata, "\n", lineparts);

	for (size_t v=0; v < lineparts.size(); v++) {
		found = lineparts[v].find(rdev);
		if (found == string::npos) {
			if (lineparts[v] != "" && lineparts[v] != "\n") {
				results.push_back(lineparts[v]);
			}
		}
	}

	return;
}

std::string get_root_device()
{
	std::string mountdata;
	std::string tmpdev;
	std::vector<string> lineparts;
	size_t found;

	mountdata = run_command("mount");
	Split(mountdata, "\n", lineparts);

	for (size_t i=0; i < lineparts.size(); i++) {
		if (lineparts[i][0] == '/') {
			found = lineparts[i].find("on /");
			if (found != string::npos) {
				tmpdev = lineparts[i].substr(0, found - 1);
			}
		}
	}

	return tmpdev;
}

std::string run_command(string cmd)
{
        string data;
        FILE *stream;
        char buffer[1024];             // Subject To Buffer Overflow Attacks

        stream = popen(cmd.c_str(), "r");
        while (fgets(buffer, 1024, stream) != NULL) {
                data.append(buffer);
        }
        pclose(stream);

        return data;
}

void handle_error(int error)
{
	return;
}

